from src.tools.chain import gen_chain
from src.tools.dms_operator import dmsoperator


def check_file_type1(file_path):
    file_text = dmsoperator.get_textfile(file_path)
    contexts = file_text
    if not contexts:
        return "其他"
    detect_prompt = f"已知以下文本是一段文本中的封面页；判断该段文本属于(生产批记录、工艺规程、工艺信息表、供应商目录、其他)中的哪一类；只用回答类别。\n===\n{contexts}\n===\n"
    res = gen_chain(detect_prompt)
    return res

def check_file_type2(file_path):
    file_text = dmsoperator.get_textfile(file_path)
    contexts = file_text
    if not contexts:
        return "其他"
    detect_prompt = f"已知以下文本是一段文本中的封面页；判断该段文本属于(普通工艺规程、多室袋工艺规程、其他)中的哪一类；只用回答类别。\n===\n{contexts}\n===\n"
    res = gen_chain(detect_prompt)
    return res


if __name__ == '__main__':
    res = check_file_type1(r"Files\202412\20241202155621pBmC7F98r342fec3ccd70c244c9918f6b47442a7ca.docx")
    print(res)
